from fastapi import APIRouter, HTTPException, Depends, Query, UploadFile, File, Form, BackgroundTasks
from typing import List, Dict, Any, Optional
from pydantic import BaseModel
import logging
from app.utils.rag_service import RAGService
from app.utils.document_processor import ProductDocumentProcessor
from app.utils.file_utils import FileManager
import os
import shutil
import uuid
import yaml
from datetime import datetime
from pathlib import Path
from fastapi.responses import JSONResponse
from app.core.config import settings
from app.utils.vectorstore_client import VectorStoreClient, DEPS_AVAILABLE
from app.utils.document_loader import load_document

# 为了向后兼容，定义VECTORSTORE_AVAILABLE
VECTORSTORE_AVAILABLE = DEPS_AVAILABLE

router = APIRouter()
logger = logging.getLogger(__name__)
rag_service = RAGService()

class QuestionRequest(BaseModel):
    question: str
    context: Optional[str] = None
    max_tokens: int = 2048
    product_ids: Optional[List[str]] = None  # 添加产品ID列表参数

class ComparisonRequest(BaseModel):
    product_ids: List[str]
    question: Optional[str] = "请分析这些产品的优势与劣势，并进行市场定位和技术参数对比。"
    max_tokens: int = 2048

class RAGInit:
    """RAG初始化和管理类"""
    
    @staticmethod
    def check_vectorstore_available() -> Dict[str, Any]:
        """检查向量存储库是否可用"""
        return {
            "available": VECTORSTORE_AVAILABLE,
            "error": None if VECTORSTORE_AVAILABLE else VectorStoreClient.get_init_error() or "向量存储库未安装"
        }
    
    @staticmethod
    async def initialize_vectorstore() -> Dict[str, Any]:
        """初始化向量存储"""
        # 首先检查依赖库是否可用
        availability = RAGInit.check_vectorstore_available()
        if not availability["available"]:
            return {
                "success": False,
                "message": f"无法初始化向量存储: {availability['error']}",
                "status": "未安装依赖"
            }
            
        try:
            # 检查向量存储目录
            vector_dir = Path(settings.VECTOR_STORE_DIR)
            os.makedirs(vector_dir, exist_ok=True)
            
            # 如果没有内容可向量化，返回等待状态
            if not os.path.exists(settings.KNOWLEDGE_BASE_DIR):
                os.makedirs(settings.KNOWLEDGE_BASE_DIR, exist_ok=True)
                return {
                    "success": True, 
                    "message": "向量存储初始化成功，但知识库为空",
                    "status": "等待内容"
                }
                
            # 加载知识库文档
            from app.utils.document_loader import load_knowledge_base
            documents = await load_knowledge_base(settings.KNOWLEDGE_BASE_DIR)
            
            if not documents:
                return {
                    "success": True,
                    "message": "向量存储初始化成功，但未找到可用文档",
                    "status": "等待内容"
                }
                
            # 分割文档
            from app.utils.document_splitter import split_documents
            chunks = await split_documents(documents)
            
            if not chunks:
                return {
                    "success": True,
                    "message": "向量存储初始化成功，但文档分割后无有效内容",
                    "status": "等待内容"
                }
                
            # 创建向量存储
            vector_store_client = VectorStoreClient.get_instance()
            success, error_msg = await vector_store_client.create_or_load_vectorstore(
                documents=chunks,
                store_path=os.path.join(vector_dir, "knowledge_base"),
                store_type="chroma"  # 使用ChromaDB存储
            )
            
            if not success:
                return {
                    "success": False,
                    "message": f"创建向量存储失败: {error_msg}",
                    "status": "初始化失败"
                }
                
            return {
                "success": True,
                "message": f"向量存储初始化成功，包含{len(chunks)}个文档块",
                "status": "运行中"
            }
            
        except Exception as e:
            logging.error(f"初始化向量存储失败: {str(e)}")
            return {
                "success": False,
                "message": f"初始化向量存储时发生错误: {str(e)}",
                "status": "初始化错误"
            }

# 定义文档相关模型
class Document(BaseModel):
    id: str
    filename: str
    file_path: str
    upload_time: str
    size: int
    type: str

@router.post("/ask", response_model=Dict[str, Any])
async def ask_question(request: QuestionRequest):
    """使用RAG回答问题"""
    try:
        # 如果提供了产品ID，则获取产品内容作为上下文
        context = request.context
        if request.product_ids and len(request.product_ids) > 0:
            logger.info(f"使用产品进行回答: {request.product_ids}")
            
            # 整合所有产品信息作为上下文
            product_contexts = []
            for product_id in request.product_ids:
                try:
                    # 获取产品详情
                    product_detail = FileManager.get_product_detail(product_id)
                    if not product_detail:
                        logger.warning(f"找不到产品: {product_id}")
                        continue
                    
                    # 产品名称
                    product_name = product_detail.get("name", product_id.split("_")[0])
                    
                    # 标记产品开始
                    product_context = f"## 产品: {product_name}\n\n"
                    
                    # 将产品详情字典序列化为YAML格式字符串
                    try:
                        yaml_content = yaml.dump(product_detail, allow_unicode=True, default_flow_style=False, sort_keys=False, Dumper=yaml.SafeDumper)
                        product_context += f"```yaml\n{yaml_content}\n```\n\n"
                    except Exception as e:
                        logger.warning(f"无法序列化产品{product_id}为YAML: {str(e)}")
                        # 作为备选，添加键值对格式
                        for key, value in product_detail.items():
                            if isinstance(value, str) or isinstance(value, (int, float, bool)):
                                product_context += f"{key}: {value}\n"
                    
                    product_contexts.append(product_context)
                except Exception as e:
                    logger.error(f"处理产品 {product_id} 时出错: {str(e)}")
            
            # 如果找到了产品上下文，添加到请求上下文中
            if product_contexts:
                product_context_text = "\n\n".join(product_contexts)
                if context:
                    context = f"{context}\n\n{product_context_text}"
                else:
                    context = product_context_text
                logger.info(f"已添加产品上下文，长度: {len(context)}")
        
        result = await RAGService.answer_question(
            query=request.question,
            context=context,
            max_tokens=request.max_tokens
        )
        return result
    except Exception as e:
        logging.error(f"问答API异常: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理问题失败: {str(e)}")

@router.post("/compare-products", response_model=Dict[str, Any])
async def compare_products(request: ComparisonRequest) -> Dict[str, Any]:
    """对比多个产品"""
    try:
        product_data = []
        
        # 获取每个产品的详细信息
        for product_id in request.product_ids:
            product_detail = FileManager.get_product_detail(product_id)
            if not product_detail:
                raise HTTPException(
                    status_code=404,
                    detail=f"找不到产品：{product_id}"
                )
            
            # 获取产品的文本提取结果
            product_texts = FileManager.get_product_image_texts(product_id)
            
            # 提取分析结果
            analysis_result = {}
            if 'analysis_result' in product_detail and product_detail['analysis_result']:
                analysis_result = product_detail['analysis_result']
            
            # 合并所有提取的文本
            extracted_text = ""
            if product_texts:
                for image_name, text_content in product_texts.items():
                    extracted_text += f"{image_name}: {text_content}\n\n"
            
            # 添加到产品数据列表
            product_data.append({
                "id": product_id,
                "name": product_detail.get("name", "未命名产品"),
                "extracted_text": extracted_text,
                "analysis_result": analysis_result
            })
        
        # 调用RAG服务进行产品对比分析
        result = await RAGService().analyze_product_comparison(
            product_data=product_data,
            max_tokens=request.max_tokens,
            question=request.question  # 传递自定义问题
        )
        
        return result
    except HTTPException as e:
        # 重新抛出HTTP异常
        raise e
    except Exception as e:
        logging.error(f"产品对比分析失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"产品对比分析失败: {str(e)}"
        )

@router.post("/vectors/initialize", response_model=Dict[str, Any])
async def initialize_vectors():
    """初始化所有产品的向量"""
    try:
        result = await ProductDocumentProcessor.process_all_products()
        return result
    except Exception as e:
        logging.error(f"初始化向量API异常: {str(e)}")
        raise HTTPException(status_code=500, detail=f"初始化向量失败: {str(e)}")

@router.post("/vectors/update/{product_id}", response_model=Dict[str, Any])
async def update_product_vectors(product_id: str):
    """更新指定产品的向量"""
    try:
        result = await ProductDocumentProcessor.update_product_vectors(product_id)
        return result
    except Exception as e:
        logging.error(f"更新产品向量API异常: {str(e)}")
        raise HTTPException(status_code=500, detail=f"更新产品向量失败: {str(e)}")

@router.delete("/vectors/{product_id}", response_model=Dict[str, Any])
async def delete_product_vectors(product_id: str):
    """删除指定产品的向量"""
    try:
        result = await ProductDocumentProcessor.delete_product_vectors(product_id)
        return result
    except Exception as e:
        logging.error(f"删除产品向量API异常: {str(e)}")
        raise HTTPException(status_code=500, detail=f"删除产品向量失败: {str(e)}")

# 添加文档上传和检索的路由
@router.post("/documents/upload", response_model=Dict[str, Any])
async def upload_documents(files: List[UploadFile] = File(...), background_tasks: BackgroundTasks = None):
    """上传文档到系统并处理"""
    result = {"success": True, "files": []}
    
    # 确保文档目录存在
    documents_dir = os.path.join("uploads", "documents")
    os.makedirs(documents_dir, exist_ok=True)
    
    # 文档索引文件路径
    index_file = os.path.join(documents_dir, "index.yaml")
    
    # 读取现有索引
    documents_index = []
    if os.path.exists(index_file):
        with open(index_file, "r", encoding="utf-8") as f:
            try:
                documents_index = yaml.safe_load(f) or []
            except Exception as e:
                logger.error(f"读取文档索引失败: {e}")
                documents_index = []
    
    for file in files:
        try:
            # 生成唯一ID和安全的文件名
            doc_id = str(uuid.uuid4())
            original_filename = file.filename
            file_extension = os.path.splitext(original_filename)[1].lower()
            
            # 检查文件类型
            allowed_extensions = [".pdf", ".docx", ".doc", ".txt"]
            if file_extension not in allowed_extensions:
                result["files"].append({
                    "filename": original_filename,
                    "status": "error",
                    "message": "不支持的文件类型"
                })
                continue
            
            # 保存文件
            safe_filename = f"{doc_id}{file_extension}"
            file_path = os.path.join(documents_dir, safe_filename)
            
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            
            # 获取文件大小
            file_size = os.path.getsize(file_path)
            
            # 添加到索引
            doc_info = {
                "id": doc_id,
                "filename": original_filename,
                "file_path": file_path,
                "upload_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "size": file_size,
                "type": file_extension[1:]  # 移除点号
            }
            
            documents_index.append(doc_info)
            
            # 将结果添加到响应
            result["files"].append({
                "id": doc_id,
                "filename": original_filename,
                "status": "success"
            })
            
            # 在后台处理文档（如索引到向量数据库）
            if background_tasks:
                background_tasks.add_task(process_document, doc_info)
            
        except Exception as e:
            logger.error(f"文档上传失败: {e}")
            result["files"].append({
                "filename": file.filename,
                "status": "error",
                "message": str(e)
            })
    
    # 保存更新的索引
    try:
        with open(index_file, "w", encoding="utf-8") as f:
            yaml.dump(documents_index, f, allow_unicode=True, Dumper=yaml.SafeDumper)
    except Exception as e:
        logger.error(f"保存文档索引失败: {e}")
        result["warning"] = "文档已上传但索引未更新"
    
    return result

@router.get("/documents/list", response_model=List[Document])
async def list_documents():
    """获取所有上传的文档列表"""
    documents_dir = os.path.join("uploads", "documents")
    index_file = os.path.join(documents_dir, "index.yaml")
    
    if not os.path.exists(index_file):
        return []
    
    try:
        with open(index_file, "r", encoding="utf-8") as f:
            documents = yaml.safe_load(f) or []
        return documents
    except Exception as e:
        logger.error(f"读取文档索引失败: {e}")
        raise HTTPException(status_code=500, detail="无法获取文档列表")

@router.delete("/documents/{doc_id}", response_model=Dict[str, Any])
async def delete_document(doc_id: str):
    """删除指定的文档"""
    documents_dir = os.path.join("uploads", "documents")
    index_file = os.path.join(documents_dir, "index.yaml")
    
    if not os.path.exists(index_file):
        raise HTTPException(status_code=404, detail="文档索引不存在")
    
    try:
        # 读取索引
        with open(index_file, "r", encoding="utf-8") as f:
            documents = yaml.safe_load(f) or []
        
        # 查找要删除的文档
        doc_to_delete = None
        updated_documents = []
        
        for doc in documents:
            if doc["id"] == doc_id:
                doc_to_delete = doc
            else:
                updated_documents.append(doc)
        
        if not doc_to_delete:
            raise HTTPException(status_code=404, detail="文档不存在")
        
        # 删除实际文件
        if os.path.exists(doc_to_delete["file_path"]):
            os.remove(doc_to_delete["file_path"])
        
        # 更新索引
        with open(index_file, "w", encoding="utf-8") as f:
            yaml.dump(updated_documents, f, allow_unicode=True, Dumper=yaml.SafeDumper)
        
        # 从向量存储中移除文档（如果适用）
        # TODO: 实现从向量存储中删除文档
        
        return {"success": True, "message": "文档已成功删除"}
    
    except Exception as e:
        logger.error(f"删除文档失败: {e}")
        raise HTTPException(status_code=500, detail=f"删除文档失败: {str(e)}")

@router.post("/documents/query", response_model=Dict[str, Any])
async def query_document(doc_id: str = Form(...), question: str = Form(...)):
    """查询特定文档内容"""
    # 1. 查找文档
    documents_dir = os.path.join("uploads", "documents")
    index_file = os.path.join(documents_dir, "index.yaml")
    
    if not os.path.exists(index_file):
        raise HTTPException(status_code=404, detail="文档索引不存在")
    
    try:
        # 读取索引
        with open(index_file, "r", encoding="utf-8") as f:
            documents = yaml.safe_load(f) or []
        
        # 查找文档
        doc_info = None
        for doc in documents:
            if doc["id"] == doc_id:
                doc_info = doc
                break
        
        if not doc_info:
            raise HTTPException(status_code=404, detail="文档不存在")
        
        # 2. 使用RAG服务查询文档
        try:
            # 创建RAG服务实例
            rag_service = RAGService()
            
            # 构建包含文档上下文的问题
            response = await rag_service.answer_document_question(
                doc_info["file_path"], 
                question
            )
            
            return {
                "success": True,
                "document": {
                    "id": doc_info["id"],
                    "filename": doc_info["filename"]
                },
                "question": question,
                "answer": response["answer"],
                "sources": response.get("sources", [])
            }
            
        except Exception as e:
            logger.error(f"文档查询失败: {e}")
            return {
                "success": False,
                "message": f"无法处理文档查询: {str(e)}",
                "fallback_answer": "抱歉，我无法处理这个文档的查询。这可能是因为文档格式不支持或系统错误。"
            }
    
    except Exception as e:
        logger.error(f"文档查询失败: {e}")
        raise HTTPException(status_code=500, detail=f"文档查询失败: {str(e)}")

# 文档处理函数
async def process_document(doc_info: Dict[str, Any]):
    """后台处理文档，例如分割和向量化"""
    try:
        logger.info(f"开始处理文档: {doc_info['filename']}")
        
        # 根据文档类型处理
        file_path = doc_info["file_path"]
        file_type = doc_info["type"]
        
        # TODO: 实现文档处理逻辑
        # 1. 读取文档内容
        # 2. 分割文档
        # 3. 存储到向量数据库
        
        logger.info(f"文档处理完成: {doc_info['filename']}")
        
    except Exception as e:
        logger.error(f"文档处理失败: {doc_info['filename']}, 错误: {e}") 